# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
RunOneStepCell
"""

from typing import Tuple, List
import torch
from torch import Tensor
from torch.nn import Module

from torch.optim import Optimizer

from .energy import WithEnergyCell
from .force import WithForceCell
from ...function.functions import get_integer, all_none, get_arguments, GLOBAL_DEVICE_SPONGE
from ...optimizer import Updater


class RunOneStepCell(Module):

    def __init__(self,
                 energy: WithEnergyCell = None,
                 force: WithForceCell = None,
                 optimizer: Optimizer = None,
                 steps: int = 1,
                 **kwargs
                 ):

        super().__init__()
        self._kwargs = get_arguments(locals(), kwargs)

        self.device = GLOBAL_DEVICE_SPONGE()

        if all_none([energy, force]):
            raise ValueError('energy and force cannot be both None!')

        self._neighbour_list_pace = None

        self.system_with_energy = energy
        if self.system_with_energy is not None:
            self.system = self.system_with_energy.system
            self.units = self.system_with_energy.units
            self._neighbour_list_pace = self.system_with_energy.neighbour_list_pace

        self.system_with_force = force
        if self.system_with_force is not None:

            force_pace = self.system_with_force.neighbour_list_pace

            if self.system_with_energy is None or self._neighbour_list_pace == 0:
                self._neighbour_list_pace = force_pace
            else:
                if force_pace not in (self._neighbour_list_pace, 0):
                    raise ValueError(
                        f'The neighbour_list_pace in WithForceCell ({force_pace}) cannot match '
                        f'the neighbour_list_pace in WithEnergyCell ({self._neighbour_list_pace}).')

        if self.system_with_energy is None:
            self.system = self.system_with_force.system

        self.optimizer = optimizer
        if self.optimizer is None:
            print('[WARNING] No optimizer! The simulation system will not be updated!')

        self.use_updater = isinstance(self.optimizer, Updater)
        self.cal_virial = True
        self.reducer_flag = False
        # self.parallel_mode = _get_parallel_mode()
        # self.reducer_flag = self.parallel_mode in (
        #     ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL)
        # if self.reducer_flag:
        #     self.mean = _get_gradients_mean()
        #     self.degree = _get_device_num()
        #     self.grad_reducer = DistributedGradReducer(
        #         self.weights, self.mean, self.degree)

        self.steps = get_integer(steps)

    @property
    def neighbour_list_pace(self) -> int:
        r"""
        update step for neighbour list.

        Return:
            int, the number of steps needed for neighbour list updating.
        """
        return self._neighbour_list_pace

    @property
    def energy_cutoff(self) -> Tensor:
        r"""
        cutoff distance for neighbour list in WithEnergyCell.

        Return:
            Tensor, cutoff distance for neighbour list in WithEnergyCell.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.cutoff

    @property
    def force_cutoff(self) -> Tensor:
        r"""
        cutoff distance for neighbour list in WithForceCell.

        Return:
            Tensor, cutoff distance for neighbour list in WithForceCell.
        """
        if self.system_with_force is None:
            return None
        return self.system_with_force.cutoff

    @property
    def length_unit(self) -> str:
        r"""
        length unit.

        Return:
            str, length unit.
        """
        return self.units.length_unit

    @property
    def energy_unit(self) -> str:
        r"""
        energy unit.

        Return:
            str, energy unit.
        """
        return self.units.energy_unit

    @property
    def num_energies(self) -> int:
        r"""
        number of energy terms :math:`U`.

        Return:
            int, number of energy terms.
        """
        if self.system_with_energy is None:
            return 0
        return self.system_with_energy.num_energies

    @property
    def energy_names(self):
        r"""
        names of energy terms.

        Return:
            list[str], names of energy terms.
        """
        if self.system_with_energy is None:
            return []
        return self.system_with_energy.energy_names

    @property
    def bias_names(self) -> List[str]:
        r"""
        name of bias potential energies.

        Return:
            list[str], the bias potential energies.
        """
        if self.system_with_energy is None:
            return []
        return self.system_with_energy.bias_names

    @property
    def num_biases(self) -> int:
        r"""
        number of bias potential energies :math:`V`.

        Return:
            int, number of bias potential energies.
        """
        if self.system_with_energy is None:
            return 0
        return self.system_with_energy.num_biases

    @property
    def energies(self) -> Tensor:
        r"""
        Tensor of potential energy components.

        Return:
            Tensor, Tensor of shape `(B, U)`. Data type is float.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.energies

    @property
    def biases(self) -> Tensor:
        r"""
        Tensor of bias potential components.

        Return:
            Tensor, Tensor of shape `(B, V)`. Data type is float.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.biases

    @property
    def bias(self) -> Tensor:
        r"""
        Tensor of the total bias potential.

        Return:
            Tensor, Tensor of shape `(B, 1)`. Data type is float.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.bias

    @property
    def bias_function(self) -> Module:
        r"""
        Cell of bias potential function.

        Return:
            Cell, bias potential function.
        """
        if self.system_with_energy is None:
            return None
        return self.system_with_energy.bias_function

    def update_neighbour_list(self):
        r"""update neighbour list."""
        if self.system_with_energy is not None:
            self.system_with_energy.update_neighbour_list()
        if self.system_with_force is not None and self.system_with_force.neighbour_list is not None:
            self.system_with_force.update_neighbour_list()
        return self

    def update_bias(self, step: int):
        r"""
        update bias potential.

        Args:
            step(int):  Simulation step to update bias potential.
        """
        if self.system_with_energy is not None:
            self.system_with_energy.update_bias(step)
        return self

    def update_wrapper(self, step: int):
        r"""
        update energy wrapper.

        Args:
            step(int):  Simulation step to update energy wrapper.
        """
        if self.system_with_energy is not None:
            self.system_with_energy.update_wrapper(step)
        return self

    def update_modifier(self, step: int):
        r"""
        update force modifier.

        Args:
            step(int):  Simulation step to update force modifier.
        """
        if self.system_with_force is not None:
            self.system_with_force.update_modifier(step)
        return self

    def set_pbc_grad(self, value: bool):
        r"""
        set whether to calculate the gradient of PBC box.

        Args:
            value(bool): Flag to judge whether to calculate the gradient of PBC box.
        """
        if self.system_with_energy is not None:
            self.system_with_energy.set_pbc_grad(value)
        if self.system_with_force is not None:
            self.system_with_force.set_pbc_grad(value)
        return self

    def set_steps(self, steps: int):
        r"""
        set steps for JIT.

        Args:
            step(int):  Simulation step for JIT.
        """
        self.steps = get_integer(steps)
        return self

    def forward(self, *inputs) -> Tuple[Tensor, Tensor]:
        r"""
        Run simulation.

        Args:
            *inputs(list): Inputs of the 'WithEnergyCell'.

        Returns:
            - energy, Tensor of shape `(B, 1)`. Data type is float. Total potential energy.
            - force, Tensor of shape `(B, A, D)`. Data type is float. Atomic force.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
            A:  Number of the atoms in the simulation system.
            D:  Spatial dimension of the simulation system. Usually is 3.
        """

        def _run_one_step(*inputs):
            r"""
            Run one step simulation.

            Args:
                *inputs(Tuple(Tensor)): Tuple of input tensors of `WithEnergyCell`.

            Returns:
                - energy, Tensor of shape `(B, 1)`. Data type is float. Total potential energy.
                - force, Tensor of shape `(B, A, D)`. Data type is float. Atomic force.

            Symbols:
                B:  Batchsize, i.e. number of walkers of the simulation.
                A:  Number of the atoms in the simulation system.
                D:  Spatial dimension of the simulation system. Usually is 3.
            """
            energy = 0
            force = 0
            virial = torch.zeros(self.system.coordinate.size(0),3,device=self.device)
            self.optimizer.zero_grad()
            if self.system_with_energy is not None:
                energy = self.system_with_energy(*inputs)
                energy.backward(torch.ones_like(energy),retain_graph=True)
                force = - self.system.coordinate.grad

                if self.system.pbc_box is not None and self.system.pbc_box.grad is not None:
                    virial = 0.5 * self.system.pbc_box.grad * self.system.pbc_box
                    self.system.pbc_box.grad.zero_()

            if self.system_with_force is not None:
                energy, force, virial = self.system_with_force(energy, force, virial)

            if self.optimizer is not None:
                if self.use_updater:
                    self.optimizer.step(energy, force, virial)
                    # if self.optimizer.step_num % 100 ==0:
                    #     print("step {0}: temperature: {1:.2f}".format(int(self.optimizer.step_num.detach()),float(temperature.cpu().detach().numpy())))
                else:
                    self.optimizer.step()
        
            return energy, force

        if self.steps == 1:
            return _run_one_step(*inputs)

        energy = None
        force = None
        for _ in range(self.steps):
            energy, force = _run_one_step(*inputs)
            self.update_neighbour_list()
            
        return energy, force
